import torch
import time
import ctypes

libsuba = ctypes.CDLL('./build/libsuba.so')

def fun2():
    time.sleep(0.001)
    libsuba.my_func(1234)

def fun1():
    fun2()
    a = torch.rand(100, 100)
    b = torch.rand(100, 100)
    c = a @ b


with torch.profiler.profile(
    activities=[torch.profiler.ProfilerActivity.CPU],  # 仅关注 CPU
    on_trace_ready=torch.profiler.tensorboard_trace_handler('./logs'),
    with_stack=True,  # 显示调用栈
) as prof:
    fun1()
libsuba.dclLlmStopProfiler()
